"""用于 LangGraph 的代码执行工具。

本模块提供了一个在沙箱环境中安全执行代码的工具，可与 LangGraph 一起使用
来执行各种编程语言的代码。
"""
import json
from typing import List, Optional
# 在执行前增加AST解析验证 AST模块仅适用于Python代码，对其他语言无效
import ast
import docker
from docker.errors import DockerException
from llm_sandbox import SandboxSession, SandboxBackend
from pydantic import BaseModel, Field
from langchain_core.tools import Tool, StructuredTool
from app.core.logging import logger
from app.core.repl.security_configuration import PYTHON_POLICY


class CodeExecutionInput(BaseModel):
    lang: str = Field(..., description="编程语言，必须是['python', 'java', 'javascript', 'cpp', 'go', 'ruby']中的一种")
    code: str = Field(..., description="要执行的代码")


def execute_code(lang: str, code: str) -> str:
    """
    在沙箱环境中执行代码。

    Args:
        lang (str): 编程语言，必须是['python', 'java', 'javascript', 'cpp', 'go', 'ruby']中的一种
        code (str): 要执行的代码

    Returns:
        str: 代码执行的输出结果
    """

    SUPPORTED_LANGS = {'PYTHON', 'JAVA', 'JAVASCRIPT', 'CPP', 'GO', 'RUBY'}
    if lang.upper() not in SUPPORTED_LANGS:
        result={"error": "不支持的语言"}
        return json.dumps(result,ensure_ascii=False)
    # 前置语法检查  这边注释掉了，感觉不是很有必要 不是很熟悉ast
    # if lang == "python":
    #     try:
    #         ast.parse(code)
    #     except SyntaxError as e:
    #         logger.warning(f"Python语法错误: {e}")
    #         result={"error": f"语法错误: {e}"}
    #         return  json.dumps(result,ensure_ascii=False)#作为工具调用结果返回给大模型

    # 安全策略 后续可以添加其他语言的
    security_policy = PYTHON_POLICY if lang == "python" else None
    # 动态生成容器ID
    container_id = f"my_{lang}_container" if lang != "python" else "my_python_container"

    try:
        client = docker.from_env()
        with SandboxSession(lang=lang,
                            backend=SandboxBackend.DOCKER,
                            verbose=True,
                            container_id=container_id,  # 目前只有 python容器
                            client=client,
                            execution_timeout=30,
                            security_policy=security_policy
                            ) as session:
            is_safe, violations = session.is_safe(code)  # 检查代码是否安全 返回值是一个元组，第一个是一个布尔值，第二个是列表，列表中的元素是违反的规则
            if not is_safe:
                # Log security violation
                violations_str = ", ".join([v.description for v in violations])
                logger.warning(f"安全检查未通过: {violations_str}")
                result={"error": f"安全检查未通过: {violations_str}"}
                return json.dumps(result, ensure_ascii=False)  # 作为工具调用结果返回给大模型

            # Execute safely
            try:
                resultOutput = session.run(code)
                result = resultOutput.stdout
                stderr = resultOutput.stderr
                if stderr:
                    return json.dumps({"error": f"执行错误: {stderr}"}, ensure_ascii=False)
                logger.info(f"成功运行代码,结果:{result}")
                result={"result": result}
                return json.dumps(result, ensure_ascii=False)
            except Exception as e:
                logger.error(f"Execution error: {e}")
                result={"error": f"执行错误: {e}"}
                return json.dumps(result, ensure_ascii=False)
    except DockerException as de:
        logger.exception(f"Docker操作失败: {de}")
        result={"error": f"Docker操作失败: {de}"}
        return json.dumps(result, ensure_ascii=False)
    except Exception as e:
        logger.exception(f"未处理的异常: {e}")
        result={"error": f"未处理的异常: {e}"}
        return json.dumps(result, ensure_ascii=False)


# 创建代码执行工具实例
execute_code_tool = StructuredTool(
    name="execute_code_tool",
    description="在安全沙箱环境中执行代码。支持多种编程语言包括PYTHON、JAVA、JAVASCRIPT等。环境配置：Python 3.12 + 全套科学计算库（NumPy, SciPy, Pandas, Matplotlib, Scikit-learn, SymPy）",
    func=execute_code,
    args_schema=CodeExecutionInput,
    handle_tool_error=True
)
